/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.zookeeper.server;

import org.apache.jute.BinaryInputArchive;
import org.apache.jute.BinaryOutputArchive;
import org.apache.jute.Record;
import org.apache.log4j.Logger;
import org.apache.zookeeper.ClientCnxn;
import org.apache.zookeeper.proto.ConnectRequest;
import org.apache.zookeeper.proto.ConnectResponse;
import org.apache.zookeeper.proto.RequestHeader;
import org.apache.zookeeper.server.paxos2.Configuration;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.SocketChannel;
import java.util.ArrayList;
import java.util.LinkedList;
import java.util.Random;
import java.util.Set;

/**
 * This class manages the socket i/o for the client. ClientCnxn maintains a list
 * of available servers to connect to and "transparently" switches servers it is
 * connected to as needed.
 */
public class ForwardingCnxn extends Thread {
    private static final Logger LOG = Logger.getLogger(ClientCnxn.class);

    private final ArrayList<InetSocketAddress> serverAddrs = new ArrayList<InetSocketAddress>();

    final ServerSession[] serverSessions;

    final Selector selector = Selector.open();

    private long clientId;

    /**
     * Set to true when close is called. Latches the connection such that we
     * don't attempt to re-connect to the server if in the middle of closing the
     * connection (client sends session disconnect to server as part of close
     * operation)
     */
    volatile boolean closing = false;

    final Object liveCountLock = new Object();
    int liveCount = 0;

    /**
     * This class allows us to pass the headers and the relevant records around.
     */
    static class Packet {

        ByteBuffer bb;

        Packet(RequestHeader header, Record record, ByteBuffer bb) {
            if (bb != null) {
                this.bb = bb;
            } else {
                try {
                    ByteArrayOutputStream baos = new ByteArrayOutputStream();
                    BinaryOutputArchive boa = BinaryOutputArchive
                            .getArchive(baos);
                    boa.writeInt(-1, "len"); // We'll fill this in later
                    header.serialize(boa, "header");
                    if (record != null) {
                        record.serialize(boa, "request");
                    }
                    baos.close();
                    this.bb = ByteBuffer.wrap(baos.toByteArray());
                    this.bb.putInt(this.bb.capacity() - 4);
                    this.bb.rewind();
                } catch (IOException e) {
                    LOG.warn("Ignoring unexpected exception", e);
                }
            }
        }

    }

    /**
     * Creates a connection object. The actual network connect doesn't get
     * established until needed. The start() instance method must be called
     * subsequent to construction.
     *
     * @throws java.io.IOException
     */
    public ForwardingCnxn(long clientId) throws IOException {

        this.clientId = clientId;
        String hostsList[] = Configuration.servers;

        for (String host : hostsList) {
            int port = 2181;
            int pidx = host.lastIndexOf(':');
            if (pidx >= 0) {
                // otherwise : is at the end of the string, ignore
                if (pidx < host.length() - 1) {
                    port = Integer.parseInt(host.substring(pidx + 1));
                }
                host = host.substring(0, pidx);
            }
            InetAddress addrs[] = InetAddress.getAllByName(host);
            for (InetAddress addr : addrs) {
                serverAddrs.add(new InetSocketAddress(addr, port));
            }
        }
        serverSessions = new ServerSession[serverAddrs.size()];
        for (int i = 0; i < serverSessions.length; i++)
            serverSessions[i] = new ServerSession(i);
    }

    public void start(){
        super.start();
        synchronized(liveCountLock){
            if(liveCount<3){
                try{
                    liveCountLock.wait(3000);
                }
                catch(InterruptedException e){
                    e.printStackTrace();
                }
            }
        }
    }


    public LinkedList<Integer> listLiveServers(){
        LinkedList<Integer> ret = new LinkedList<Integer>();
        for(int i=0; i<serverSessions.length; i++){
            if(serverSessions[i] != null && serverSessions[i].sockKey != null && serverSessions[i].sockKey.isValid())
                ret.add(i+1);
        }
        return ret;
    }

    public boolean isServerLive(int id){
        return serverSessions[id] != null && serverSessions[id].sockKey != null && serverSessions[id].sockKey.isValid()
                && serverSessions[id].initialized;
    }

    Object eventOfDeath = new Object();

    final static UncaughtExceptionHandler uncaughtExceptionHandler = new UncaughtExceptionHandler() {
        public void uncaughtException(Thread t, Throwable e) {
            e.printStackTrace();
            LOG.error("from " + t.getName(), e);
        }
    };

    /**
     * This class services the outgoing request queue and generates the heart
     * beats. It also spawns the ReadThread.
     */
    class ServerSession {
        SelectionKey sockKey;

        final ByteBuffer lenBuffer = ByteBuffer.allocateDirect(4);

        ByteBuffer incomingBuffer = lenBuffer;

        boolean initialized;


        long sentCount = 0;
        long recvCount = 0;

        int id;

        /**
         * These are the packets that need to be sent.
         */
        private final LinkedList<Packet> outgoingQueue = new LinkedList<Packet>();


        void readLength() throws IOException {
            int len = incomingBuffer.getInt();
            if (len < 0) {
                throw new IOException("Packet len" + len + " is out of range!");
            }
            incomingBuffer = ByteBuffer.allocate(len);
        }

        void readConnectResult() throws IOException {
            ByteBufferInputStream bbis = new ByteBufferInputStream(
                    incomingBuffer);
            BinaryInputArchive bbia = BinaryInputArchive.getArchive(bbis);
            ConnectResponse conRsp = new ConnectResponse();
            conRsp.deserialize(bbia, "connect");
            LOG.info("Session establishment complete on server "
                    + ((SocketChannel) sockKey.channel()).socket()
                    .getRemoteSocketAddress());
        }

        void readResponse() throws IOException {
            throw new IOException("This should not happen");
        }

        /**
         * @return true if a packet was received
         * @throws InterruptedException
         * @throws java.io.IOException
         */
        boolean doIO() throws InterruptedException, IOException {
            boolean packetReceived = false;
	    if(sockKey == null || !sockKey.isValid())
		return false;
            SocketChannel sock = (SocketChannel) sockKey.channel();
	    //System.out.println("doIO");
            if (sock == null) {
                throw new IOException("Socket is null!");
            }
            if (sockKey.isReadable()) {
                int rc = sock.read(incomingBuffer);
                if (rc < 0) {
                    throw new IOException(
                            "Unable to read additional data from server likely server has closed socket");
                    //cleanup();

                    //return false;
                }
                if (!incomingBuffer.hasRemaining()) {
                    incomingBuffer.flip();
                    if (incomingBuffer == lenBuffer) {
                        recvCount++;
                        readLength();
                    } else if (!initialized) {
                        readConnectResult();
                        enableRead();
                        //if (!outgoingQueue.isEmpty()) {
                            enableWrite();
                        //}
                        lenBuffer.clear();
                        incomingBuffer = lenBuffer;
                        packetReceived = true;
                        initialized = true;
                        synchronized(liveCountLock){
                            liveCount++;
			    if(liveCount==3)
                                liveCountLock.notifyAll();
                        }
                    } else {
                        readResponse();
                        lenBuffer.clear();
                        incomingBuffer = lenBuffer;
                        packetReceived = true;
                    }
                }
            }
            if (sockKey.isWritable()) {
                synchronized (outgoingQueue) {
                    if (!outgoingQueue.isEmpty()) {
                        ByteBuffer pbb = outgoingQueue.getFirst().bb;
			/*if(pbb.position()==0){
			    ByteBuffer len = ByteBuffer.allocate(4);
			    len.putInt(pbb.capacity());
			    System.out.println("len="+pbb.capacity());
			    while(len.hasRemaining())
			    	sock.write(len);
			}*/
                        sock.write(pbb);
			//System.out.println("sock write");
                        if (!pbb.hasRemaining()) {
                            sentCount++;
                            Packet p = outgoingQueue.removeFirst();
                        }
                    }
                }
            }
	    synchronized(outgoingQueue){
            if (outgoingQueue.isEmpty()) {
                disableWrite();
            } else {
                enableWrite();
            }
	    }
            return packetReceived;
        }

        synchronized private void enableWrite() {
            if (sockKey != null && sockKey.isValid()) {
                int i = sockKey.interestOps();
                if ((i & SelectionKey.OP_WRITE) == 0) {
                    sockKey.interestOps(i | SelectionKey.OP_WRITE);
                }
            }
        }

        synchronized private void disableWrite() {
            if (sockKey != null && sockKey.isValid()) {
                int i = sockKey.interestOps();
                if ((i & SelectionKey.OP_WRITE) != 0) {
                    sockKey.interestOps(i & (~SelectionKey.OP_WRITE));
                }
            }
        }

        synchronized private void enableRead() {
            if (sockKey != null && sockKey.isValid()) {
                int i = sockKey.interestOps();
                if ((i & SelectionKey.OP_READ) == 0) {
                    sockKey.interestOps(i | SelectionKey.OP_READ);
                }
            }
        }

        synchronized private void disableRead() {
            if (sockKey != null && sockKey.isValid()) {
                int i = sockKey.interestOps();
                if ((i & SelectionKey.OP_READ) != 0) {
                    sockKey.interestOps(i & (~SelectionKey.OP_READ));
                }
            }
        }

        ServerSession(int id) {
            this.id = id;
            setUncaughtExceptionHandler(uncaughtExceptionHandler);
            setDaemon(true);
        }

        private void primeConnection(SelectionKey k) throws IOException {
            LOG.info("Socket connection established to "
                    + ((SocketChannel) sockKey.channel()).socket()
                    .getRemoteSocketAddress() + ", initiating session");
            ConnectRequest conReq = new ConnectRequest(0, 0, clientId);
            ByteArrayOutputStream baos = new ByteArrayOutputStream();
            BinaryOutputArchive boa = BinaryOutputArchive.getArchive(baos);
            boa.writeInt(-1, "len");
            conReq.serialize(boa, "connect");
            baos.close();
            ByteBuffer bb = ByteBuffer.wrap(baos.toByteArray());
            bb.putInt(bb.capacity() - 4);
            bb.rewind();
            synchronized (outgoingQueue) {
                outgoingQueue.addFirst((new Packet(null, null, bb)));
            }
            synchronized (this) {
                k.interestOps(SelectionKey.OP_READ | SelectionKey.OP_WRITE);
            }
            if (LOG.isDebugEnabled()) {
                LOG.debug("Session establishment request sent on "
                        + ((SocketChannel) sockKey.channel()).socket()
                        .getRemoteSocketAddress());
            }
        }

        Random r = new Random(System.nanoTime());

        private long lastConnectTime = -1;
        private void startConnect() throws IOException {

            /*try {
                Thread.sleep(r.nextInt(1000));
            } catch (InterruptedException e1) {
                LOG.warn("Unexpected exception", e1);
            } */

            if(lastConnectTime != -1 && System.currentTimeMillis() - lastConnectTime < 5000)
                return;
            lastConnectTime = System.currentTimeMillis();

            InetSocketAddress addr = serverAddrs.get(id);
            LOG.info("Opening socket connection to server " + addr);
            SocketChannel sock;
            sock = SocketChannel.open();
            sock.configureBlocking(false);
            sock.socket().setSoLinger(false, -1);
            sock.socket().setTcpNoDelay(true);
            setName(getName().replaceAll("\\(.*\\)",
                    "(" + addr.getHostName() + ":" + addr.getPort() + ")"));
            sockKey = sock.register(selector, SelectionKey.OP_CONNECT);
            sockKey.attach(this);
            if (sock.connect(addr)) {
                primeConnection(sockKey);
            }
            initialized = false;

            /*
                * Reset incomingBuffer
                */
            lenBuffer.clear();
            incomingBuffer = lenBuffer;
        }

        private void cleanup() {
            if (sockKey != null) {
                SocketChannel sock = (SocketChannel) sockKey.channel();
		System.out.println("cleanup "+sockKey+" "+sock);
                sockKey.cancel();
                try {
                    sock.socket().shutdownInput();
                } catch (IOException e) {
                    if (LOG.isDebugEnabled()) {
                        LOG
                                .debug(
                                        "Ignoring exception during shutdown input",
                                        e);
                    }
                }
                try {
                    sock.socket().shutdownOutput();
                } catch (IOException e) {
                    if (LOG.isDebugEnabled()) {
                        LOG.debug("Ignoring exception during shutdown output",
                                e);
                    }
                }
                try {
                    sock.socket().close();
                } catch (IOException e) {
                    if (LOG.isDebugEnabled()) {
                        LOG.debug("Ignoring exception during socket close", e);
                    }
                }
                try {
                    sock.close();
                } catch (IOException e) {
                    if (LOG.isDebugEnabled()) {
                        LOG.debug("Ignoring exception during channel close", e);
                    }
                }
            }
            try {
                Thread.sleep(100);
            } catch (InterruptedException e) {
                if (LOG.isDebugEnabled()) {
                    LOG
                            .debug("ServerSession interrupted during sleep, ignoring");
                }
            }
            sockKey = null;
            synchronized (outgoingQueue) {
                outgoingQueue.clear();
            }
        }
    }

    private static final String RETRY_CONN_MSG = ", closing socket connection and attempting reconnect";

    @Override
    public void run() {
        long now = System.currentTimeMillis();
        long lastHeard = now;
        long lastSend = now;
        ServerSession ss = null;
        while (true) {
            try {
                for (int i = 0; i < serverSessions.length; i++) {
                    if (serverSessions[i].sockKey == null) {
                        serverSessions[i].startConnect();
                    }
                }
		/*if (serverSessions[serverId].sockKey == null) {
                        serverSessions[serverId].startConnect();
                }*/

		//System.out.println("before select "+System.currentTimeMillis());
                selector.select(1000);
                Set<SelectionKey> selected = null;
                synchronized (this) {
                    selected = selector.selectedKeys();
                }
		//System.out.println("after select "+selected.size()+" "+System.currentTimeMillis());
		//System.out.println("select size "+selected.size());
                // Everything below and until we get back to the select is
                // non blocking, so time is effectively a constant. That is
                // Why we just have to do this once, here
                now = System.currentTimeMillis();
                for (SelectionKey k : selected) {
		//    System.out.println(k);
                    SocketChannel sc = ((SocketChannel) k.channel());
                    ss = (ServerSession) k.attachment();
                    if ((k.readyOps() & SelectionKey.OP_CONNECT) != 0) {
                        if (sc.finishConnect()) {
                            lastHeard = now;
                            lastSend = now;
                            ss.primeConnection(k);
                        }
                    } else if ((k.readyOps() & (SelectionKey.OP_READ | SelectionKey.OP_WRITE)) != 0) {
                        if (ss.outgoingQueue.size() > 0) {
                            // We have something to send so it's the same
                            // as if we do the send now.
                            lastSend = now;
                        }
                        if (ss.doIO()) {
                            lastHeard = now;
                        }
                    }
                    if (ss.outgoingQueue.size() > 0) {
                        ss.enableWrite();
                    } else {
                        //ss.disableWrite();
                    }
                }
                selected.clear();
            } catch (Exception e) {
                if (closing) {
                    if (LOG.isDebugEnabled()) {
                        // closing so this is expected
                        LOG
                                .debug("An exception was thrown while closing send thread for session 0x : "
                                        + e.getMessage());
                    }
                    break;
                } else {

                    //e.printStackTrace();
                    LOG.warn("Got exception", e);
                    if (ss != null)
                        ss.cleanup();

                    now = System.currentTimeMillis();
                    lastHeard = now;
                    lastSend = now;
                }
            }
        }
        for (int i = 0; i < serverSessions.length; i++)
            serverSessions[i].cleanup();
        try {
            selector.close();
        } catch (IOException e) {
            LOG.warn("Ignoring exception during selector close", e);
        }
        ZooTrace.logTraceMessage(LOG, ZooTrace.getTextTraceLevel(),
                "ServerSession exitedloop.");
    }

    /*
      * public void close() { synchronized (this) { selector.wakeup(); } }
      */

    /**
     * Shutdown the send/event threads. This method should not be called
     * directly - rather it should be called as part of close operation. This
     * method is primarily here to allow the tests to verify disconnection
     * behavior.
     */
    public void disconnect() {
        if (LOG.isDebugEnabled()) {
            LOG.debug("Disconnecting client ");
        }

        // sendThread.close();

        synchronized (this) {
            selector.wakeup();
        }
    }

    /**
     * Close the connection, which includes; send session disconnect to the
     * server, shutdown the send/event threads.
     *
     * @throws java.io.IOException
     */
    public void close() throws IOException {
        if (LOG.isDebugEnabled()) {
            LOG.debug("Closing client for session: ");
        }


        disconnect();

    }


    Packet queuePacket(ByteBuffer bb, int sid) {
	//if(sid!=serverId)
	//    throw new RuntimeException("sid ="+sid+" serverId="+serverId);
	//System.out.println("queue to "+sid+" "+System.currentTimeMillis());
        Packet packet = null;
        ServerSession server = serverSessions[sid];
        synchronized (server.outgoingQueue) {
            ByteBuffer len = ByteBuffer.allocate(4);
            len.putInt(bb.capacity());
	    len.rewind();
            //System.out.println("len="+bb.capacity());
	    server.outgoingQueue.add(new Packet(null, null, len));

            packet = new Packet(null, null, bb);

            server.outgoingQueue.add(packet);
	    //if(server.outgoingQueue.size()>100)
	//	System.out.println("forwarder "+server.outgoingQueue.size());
            server.enableWrite();
        }
        synchronized (this) {
            selector.wakeup();
		}
	//System.out.println("end queue to "+sid+" "+System.currentTimeMillis());
		return packet;
	}

}
